{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "7f4e4862",
   "metadata": {},
   "outputs": [],
   "source": [
    "zidian = {\n",
    "    '<PAD>': 0,\n",
    "    '1': 1,\n",
    "    '2': 2,\n",
    "    '3': 3,\n",
    "    '4': 4,\n",
    "    '5': 5,\n",
    "    '6': 6,\n",
    "    '7': 7,\n",
    "    '8': 8,\n",
    "    '9': 9,\n",
    "    '0': 10,\n",
    "    'Jan': 11,\n",
    "    'Feb': 12,\n",
    "    'Mar': 13,\n",
    "    'Apr': 14,\n",
    "    'May': 15,\n",
    "    'Jun': 16,\n",
    "    'Jul': 17,\n",
    "    'Aug': 18,\n",
    "    'Sep': 19,\n",
    "    'Oct': 20,\n",
    "    'Nov': 21,\n",
    "    'Dec': 22,\n",
    "    '-': 23,\n",
    "    '/': 24,\n",
    "    '<SOS>': 25,\n",
    "    '<EOS>': 26,\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "652b0528",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 7,  7, 23,  1,  1, 23, 10,  5],\n",
       "         [ 8, 10, 23,  1,  2, 23,  1, 10],\n",
       "         [ 2,  4, 23, 10,  9, 23, 10,  4],\n",
       "         [ 9,  6, 23, 10,  4, 23,  2,  8],\n",
       "         [10,  5, 23,  1,  2, 23,  2,  5]]),\n",
       " torch.Size([100, 8]),\n",
       " tensor([[25, 10,  5, 24, 21, 24,  1,  9,  7,  7, 26],\n",
       "         [25,  1, 10, 24, 22, 24,  1,  9,  8, 10, 26],\n",
       "         [25, 10,  4, 24, 19, 24,  2, 10,  2,  4, 26],\n",
       "         [25,  2,  8, 24, 14, 24,  1,  9,  9,  6, 26],\n",
       "         [25,  2,  5, 24, 22, 24,  2, 10, 10,  5, 26]]),\n",
       " torch.Size([100, 11]))"
      ]
     },
     "execution_count": 2,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "import numpy as np\n",
    "import datetime\n",
    "from torch.utils.data import Dataset, DataLoader\n",
    "\n",
    "\n",
    "class DateDataset(Dataset):\n",
    "    def __init__(self):\n",
    "        pass\n",
    "\n",
    "    def __len__(self):\n",
    "        return 2000\n",
    "\n",
    "    def __getitem__(self, index):\n",
    "        #随机生成一个日期\n",
    "        date = np.random.randint(143835585, 2043835585)\n",
    "        date = datetime.datetime.fromtimestamp(date)\n",
    "\n",
    "        #格式化成两种格式\n",
    "        #05-06-15\n",
    "        #15/Jun/2005\n",
    "        date_cn = date.strftime(\"%y-%m-%d\")\n",
    "        date_en = date.strftime(\"%d/%b/%Y\")\n",
    "\n",
    "        #中文的就是简单的拿字典编码就行了\n",
    "        date_cn_code = [zidian[v] for v in date_cn]\n",
    "\n",
    "        #英文的,首先要在收尾加上标志位,然后用字典编码\n",
    "        date_en_code = []\n",
    "        date_en_code += [zidian['<SOS>']]\n",
    "        date_en_code += [zidian[v] for v in date_en[:3]]\n",
    "        date_en_code += [zidian[date_en[3:6]]]\n",
    "        date_en_code += [zidian[v] for v in date_en[6:]]\n",
    "        date_en_code += [zidian['<EOS>']]\n",
    "\n",
    "        return torch.LongTensor(date_cn_code), torch.LongTensor(date_en_code)\n",
    "\n",
    "\n",
    "dataloader = DataLoader(dataset=DateDataset(),\n",
    "                        batch_size=100,\n",
    "                        shuffle=True,\n",
    "                        drop_last=True)\n",
    "\n",
    "#遍历数据\n",
    "for i, data in enumerate(dataloader):\n",
    "    sample = data\n",
    "    break\n",
    "    \n",
    "sample[0][:5], sample[0].shape, sample[1][:5], sample[1].shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "5b4cd93c",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 0.0302,  0.0167,  0.0792,  0.1524,  0.0622,  0.0292,  0.0118,  0.2101,\n",
       "           0.0512, -0.1322, -0.1615,  0.1481,  0.0475, -0.1091,  0.1383,  0.0297,\n",
       "          -0.0067,  0.0993, -0.1521, -0.0160,  0.0470, -0.1038,  0.0010,  0.0759,\n",
       "          -0.0647,  0.0488, -0.0588],\n",
       "         [ 0.0115, -0.0475,  0.0737,  0.1609,  0.0372,  0.0422, -0.0366,  0.2170,\n",
       "           0.0935, -0.1141, -0.1095,  0.1391,  0.0552, -0.1908,  0.1464,  0.1299,\n",
       "           0.0055,  0.1804, -0.1547, -0.0092,  0.0187, -0.1095,  0.0392,  0.0814,\n",
       "          -0.0992, -0.0047, -0.0558]], grad_fn=<SliceBackward>),\n",
       " torch.Size([100, 10, 27]))"
      ]
     },
     "execution_count": 3,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "class Attention(nn.Module):\n",
    "    def __init__(self):\n",
    "        super().__init__()\n",
    "\n",
    "        #encoder\n",
    "        #一共27个词,编码成16维向量\n",
    "        self.encoder_embed = nn.Embedding(num_embeddings=27, embedding_dim=16)\n",
    "\n",
    "        #输入是16维向量,隐藏层是32维向量\n",
    "        self.encoder = nn.LSTM(input_size=16,\n",
    "                               hidden_size=32,\n",
    "                               num_layers=1,\n",
    "                               batch_first=True)\n",
    "\n",
    "        #decoder\n",
    "        #一共27个词,编码成16维向量\n",
    "        self.decoder_embed = nn.Embedding(num_embeddings=27, embedding_dim=16)\n",
    "\n",
    "        #输入是16维向量,隐藏层是32维向量\n",
    "        self.decoder_cell = nn.LSTMCell(input_size=16, hidden_size=32)\n",
    "\n",
    "        #输入是64维向量,输出是27分类\n",
    "        self.out_fc = nn.Linear(in_features=64, out_features=27)\n",
    "\n",
    "        #注意力全连接层\n",
    "        self.attn_fc = nn.Linear(in_features=32, out_features=32)\n",
    "\n",
    "    def get_attn(self, out_x, h):\n",
    "\n",
    "        #[b,32] -> [b,32]\n",
    "        attn = self.attn_fc(h)\n",
    "\n",
    "        #[b,32] -> [b,1,32]\n",
    "        attn = attn.unsqueeze(dim=1)\n",
    "\n",
    "        #[b,8,32] -> [b,32,8]\n",
    "        out_x_T = out_x.permute(0, 2, 1)\n",
    "\n",
    "        #[b,1,32],[b,32,8] -> [b,1,8]\n",
    "        attn = torch.matmul(attn, out_x_T)\n",
    "\n",
    "        #[b,1,8] -> [b,1,8]\n",
    "        attn = F.softmax(attn, dim=2)\n",
    "\n",
    "        #[b,1,8],[b,8,32] -> [b,1,32]\n",
    "        attn = torch.matmul(attn, out_x)\n",
    "\n",
    "        #[b,1,32] -> [b,32]\n",
    "        attn = attn.squeeze()\n",
    "\n",
    "        return attn\n",
    "\n",
    "    def forward(self, x, y):\n",
    "\n",
    "        #x编码\n",
    "        #[b,8] -> [b,8,16]\n",
    "        x = self.encoder_embed(x)\n",
    "\n",
    "        #进入循环网络,得到记忆\n",
    "        #[b,8,16] -> [b,8,32],[1,b,32],[1,b,32]\n",
    "        out_x, (h, c) = self.encoder(x, None)\n",
    "\n",
    "        #[1,b,32],[1,b,32] -> [b,32],[b,32]\n",
    "        h = h.squeeze()\n",
    "        c = c.squeeze()\n",
    "\n",
    "        #丢弃y的最后一个词\n",
    "        #因为训练的时候是以y的每一个词输入,预测下一个词\n",
    "        #所以不需要最后一个词\n",
    "        #[b,11] -> [b,10]\n",
    "        y = y[:, :-1]\n",
    "\n",
    "        #y编码\n",
    "        #[b,10] -> [b,10,16]\n",
    "        y = self.decoder_embed(y)\n",
    "\n",
    "        #用cell遍历y的每一个词\n",
    "        outs = []\n",
    "        for i in range(10):\n",
    "\n",
    "            attn = self.get_attn(out_x, h)\n",
    "\n",
    "            #把y的每个词依次输入循环网络\n",
    "            #第一个词的记忆是x的最后一个词的记忆\n",
    "            #往后每个词的记忆是上一个词的记忆\n",
    "            #[b,16] -> [b,32],[b,32]\n",
    "            h, c = self.decoder_cell(y[:, i], (h, c))\n",
    "\n",
    "            #[b,32],[b,32] -> [b,64]\n",
    "            attn = torch.cat([attn, h], dim=1)\n",
    "\n",
    "            #把每一步的记忆输出成词\n",
    "            #[b,64] -> [b,27]\n",
    "            out = self.out_fc(attn)\n",
    "            outs.append(out)\n",
    "\n",
    "        #把所有的输出词组合成一句话\n",
    "        outs = torch.stack(outs, dim=0)\n",
    "        #[10,b,27] -> #[b,10,27]\n",
    "        outs = outs.permute(1, 0, 2)\n",
    "\n",
    "        return outs\n",
    "\n",
    "\n",
    "model = Attention()\n",
    "\n",
    "out = model(sample[0], sample[1])\n",
    "out[0, :2], out.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "20f541cd",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "0 1.8220133781433105\n",
      "10 0.01074887439608574\n",
      "20 0.0023338994942605495\n",
      "30 0.0011760876514017582\n",
      "40 0.0007431698031723499\n",
      "50 0.0004605210560839623\n",
      "60 0.0003044096229132265\n",
      "70 0.0002333312586415559\n",
      "80 0.00019364364561624825\n",
      "90 0.00015617988537997007\n"
     ]
    }
   ],
   "source": [
    "loss_func = nn.CrossEntropyLoss()\n",
    "optimizer = torch.optim.Adam(model.parameters(), lr=1e-2)\n",
    "\n",
    "model.train()\n",
    "for epoch in range(100):\n",
    "    for i, data in enumerate(dataloader):\n",
    "        x, y = data\n",
    "\n",
    "        optimizer.zero_grad()\n",
    "\n",
    "        #计算输出\n",
    "        y_pred = model(x, y)\n",
    "\n",
    "        #丢弃y的第一个词\n",
    "        #因为训练的时候是以y的每一个词输入,预测下一个词\n",
    "        #所以在计算loss的时候不需要第一个词\n",
    "        #[b,11] -> [b,10]\n",
    "        y = y[:, 1:]\n",
    "\n",
    "        #打平,不然计算不了loss\n",
    "        #[b,10,27] -> [b*10,27]\n",
    "        y_pred = y_pred.reshape(-1, 27)\n",
    "\n",
    "        #[b,10] -> [b*10]\n",
    "        y = y.reshape(-1)\n",
    "\n",
    "        loss = loss_func(y_pred, y)\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "\n",
    "    if epoch % 10 == 0:\n",
    "        print(epoch, loss.item())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "06874eb1",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "('77-11-05', '<SOS>05/Nov/1977<EOS>')"
      ]
     },
     "execution_count": 5,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "#构造反转的字典\n",
    "reverse_zidian = {}\n",
    "for k, v in zidian.items():\n",
    "    reverse_zidian[v] = k\n",
    "reverse_zidian\n",
    "\n",
    "\n",
    "#数字化的句子转字符串\n",
    "def seq_to_str(seq):\n",
    "    seq = seq.detach().numpy()\n",
    "    return ''.join([reverse_zidian[idx] for idx in seq])\n",
    "\n",
    "\n",
    "seq_to_str(sample[0][0]), seq_to_str(sample[1][0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "af32f9b3",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "16-08-25 <SOS>25/Aug/2016<EOS> 25/Aug/2016\n",
      "32-01-09 <SOS>09/Jan/2032<EOS> 09/Jan/2032\n",
      "07-02-08 <SOS>08/Feb/2007<EOS> 08/Feb/2007\n",
      "90-02-27 <SOS>27/Feb/1990<EOS> 27/Feb/1990\n",
      "27-08-07 <SOS>07/Aug/2027<EOS> 07/Aug/2027\n",
      "90-07-09 <SOS>09/Jul/1990<EOS> 09/Jul/1990\n",
      "21-07-22 <SOS>22/Jul/2021<EOS> 22/Jul/2021\n",
      "84-02-28 <SOS>28/Feb/1984<EOS> 28/Feb/1984\n",
      "30-05-25 <SOS>25/May/2030<EOS> 25/May/2030\n",
      "87-05-30 <SOS>30/May/1987<EOS> 30/May/1987\n",
      "07-02-12 <SOS>12/Feb/2007<EOS> 12/Feb/2007\n",
      "05-01-24 <SOS>24/Jan/2005<EOS> 24/Jan/2005\n",
      "26-07-21 <SOS>21/Jul/2026<EOS> 21/Jul/2026\n",
      "28-10-09 <SOS>09/Oct/2028<EOS> 09/Oct/2028\n",
      "17-06-03 <SOS>03/Jun/2017<EOS> 03/Jun/2017\n",
      "05-11-09 <SOS>09/Nov/2005<EOS> 09/Nov/2005\n",
      "95-09-13 <SOS>13/Sep/1995<EOS> 13/Sep/1995\n",
      "87-06-07 <SOS>07/Jun/1987<EOS> 07/Jun/1987\n",
      "89-11-13 <SOS>13/Nov/1989<EOS> 13/Nov/1989\n",
      "31-11-14 <SOS>14/Nov/2031<EOS> 14/Nov/2031\n",
      "91-11-05 <SOS>05/Nov/1991<EOS> 05/Nov/1991\n",
      "98-06-05 <SOS>05/Jun/1998<EOS> 05/Jun/1998\n",
      "07-01-29 <SOS>29/Jan/2007<EOS> 29/Jan/2007\n",
      "79-08-22 <SOS>22/Aug/1979<EOS> 22/Aug/1979\n",
      "15-12-16 <SOS>16/Dec/2015<EOS> 16/Dec/2015\n",
      "98-12-19 <SOS>19/Dec/1998<EOS> 19/Dec/1998\n",
      "84-05-12 <SOS>12/May/1984<EOS> 12/May/1984\n",
      "30-07-04 <SOS>04/Jul/2030<EOS> 04/Jul/2030\n",
      "00-07-20 <SOS>20/Jul/2000<EOS> 20/Jul/2000\n",
      "01-04-14 <SOS>14/Apr/2001<EOS> 14/Apr/2001\n",
      "04-07-06 <SOS>06/Jul/2004<EOS> 06/Jul/2004\n",
      "81-05-09 <SOS>09/May/1981<EOS> 09/May/1981\n",
      "91-02-01 <SOS>01/Feb/1991<EOS> 01/Feb/1991\n",
      "16-12-23 <SOS>23/Dec/2016<EOS> 23/Dec/2016\n",
      "91-06-08 <SOS>08/Jun/1991<EOS> 08/Jun/1991\n",
      "07-04-21 <SOS>21/Apr/2007<EOS> 21/Apr/2007\n",
      "22-10-24 <SOS>24/Oct/2022<EOS> 24/Oct/2022\n",
      "26-04-22 <SOS>22/Apr/2026<EOS> 22/Apr/2026\n",
      "31-02-23 <SOS>23/Feb/2031<EOS> 23/Feb/2031\n",
      "31-05-15 <SOS>15/May/2031<EOS> 15/May/2031\n",
      "22-10-05 <SOS>05/Oct/2022<EOS> 05/Oct/2022\n",
      "96-06-18 <SOS>18/Jun/1996<EOS> 18/Jun/1996\n",
      "79-03-26 <SOS>26/Mar/1979<EOS> 26/Mar/1979\n",
      "82-05-30 <SOS>30/May/1982<EOS> 30/May/1982\n",
      "27-11-19 <SOS>19/Nov/2027<EOS> 19/Nov/2027\n",
      "19-01-02 <SOS>02/Jan/2019<EOS> 02/Jan/2019\n",
      "93-07-07 <SOS>07/Jul/1993<EOS> 07/Jul/1993\n",
      "82-09-08 <SOS>08/Sep/1982<EOS> 08/Sep/1982\n",
      "83-10-02 <SOS>02/Oct/1983<EOS> 02/Oct/1983\n",
      "24-08-18 <SOS>18/Aug/2024<EOS> 18/Aug/2024\n",
      "05-05-28 <SOS>28/May/2005<EOS> 28/May/2005\n",
      "12-11-26 <SOS>26/Nov/2012<EOS> 26/Nov/2012\n",
      "08-04-25 <SOS>25/Apr/2008<EOS> 25/Apr/2008\n",
      "27-05-29 <SOS>29/May/2027<EOS> 29/May/2027\n",
      "84-02-26 <SOS>26/Feb/1984<EOS> 26/Feb/1984\n",
      "20-10-14 <SOS>14/Oct/2020<EOS> 14/Oct/2020\n",
      "74-09-13 <SOS>13/Sep/1974<EOS> 13/Sep/1974\n",
      "83-07-30 <SOS>30/Jul/1983<EOS> 30/Jul/1983\n",
      "14-08-25 <SOS>25/Aug/2014<EOS> 25/Aug/2014\n",
      "06-11-13 <SOS>13/Nov/2006<EOS> 13/Nov/2006\n",
      "13-03-04 <SOS>04/Mar/2013<EOS> 04/Mar/2013\n",
      "75-07-19 <SOS>19/Jul/1975<EOS> 19/Jul/1975\n",
      "93-12-24 <SOS>24/Dec/1993<EOS> 24/Dec/1993\n",
      "03-12-03 <SOS>03/Dec/2003<EOS> 03/Dec/2003\n",
      "19-01-28 <SOS>28/Jan/2019<EOS> 28/Jan/2019\n",
      "16-05-18 <SOS>18/May/2016<EOS> 18/May/2016\n",
      "01-04-27 <SOS>27/Apr/2001<EOS> 27/Apr/2001\n",
      "96-10-09 <SOS>09/Oct/1996<EOS> 09/Oct/1996\n",
      "98-10-22 <SOS>22/Oct/1998<EOS> 22/Oct/1998\n",
      "96-04-23 <SOS>23/Apr/1996<EOS> 23/Apr/1996\n",
      "14-04-25 <SOS>25/Apr/2014<EOS> 25/Apr/2014\n",
      "79-07-02 <SOS>02/Jul/1979<EOS> 02/Jul/1979\n",
      "31-05-20 <SOS>20/May/2031<EOS> 20/May/2031\n",
      "07-02-06 <SOS>06/Feb/2007<EOS> 06/Feb/2007\n",
      "01-01-31 <SOS>31/Jan/2001<EOS> 31/Jan/2001\n",
      "12-08-12 <SOS>12/Aug/2012<EOS> 12/Aug/2012\n",
      "79-03-12 <SOS>12/Mar/1979<EOS> 12/Mar/1979\n",
      "15-02-15 <SOS>15/Feb/2015<EOS> 15/Feb/2015\n",
      "33-10-23 <SOS>23/Oct/2033<EOS> 23/Oct/2033\n",
      "78-08-11 <SOS>11/Aug/1978<EOS> 11/Aug/1978\n",
      "21-12-31 <SOS>31/Dec/2021<EOS> 31/Dec/2021\n",
      "04-11-15 <SOS>15/Nov/2004<EOS> 15/Nov/2004\n",
      "04-07-30 <SOS>30/Jul/2004<EOS> 30/Jul/2004\n",
      "76-10-25 <SOS>25/Oct/1976<EOS> 25/Oct/1976\n",
      "08-06-30 <SOS>30/Jun/2008<EOS> 30/Jun/2008\n",
      "26-12-02 <SOS>02/Dec/2026<EOS> 02/Dec/2026\n",
      "17-03-17 <SOS>17/Mar/2017<EOS> 17/Mar/2017\n",
      "09-11-13 <SOS>13/Nov/2009<EOS> 13/Nov/2009\n",
      "30-11-30 <SOS>30/Nov/2030<EOS> 30/Nov/2030\n",
      "78-08-16 <SOS>16/Aug/1978<EOS> 16/Aug/1978\n",
      "19-11-17 <SOS>17/Nov/2019<EOS> 17/Nov/2019\n",
      "87-06-02 <SOS>02/Jun/1987<EOS> 02/Jun/1987\n",
      "21-10-11 <SOS>11/Oct/2021<EOS> 11/Oct/2021\n",
      "24-03-06 <SOS>06/Mar/2024<EOS> 06/Mar/2024\n",
      "06-03-26 <SOS>26/Mar/2006<EOS> 26/Mar/2006\n",
      "30-01-07 <SOS>07/Jan/2030<EOS> 07/Jan/2030\n",
      "05-01-14 <SOS>14/Jan/2005<EOS> 14/Jan/2005\n",
      "84-12-29 <SOS>29/Dec/1984<EOS> 29/Dec/1984\n",
      "22-12-23 <SOS>23/Dec/2022<EOS> 23/Dec/2022\n",
      "11-02-08 <SOS>08/Feb/2011<EOS> 08/Feb/2011\n"
     ]
    }
   ],
   "source": [
    "#预测\n",
    "def predict(x):\n",
    "    model.eval()\n",
    "\n",
    "    #x编码\n",
    "    #[b,8] -> [b,8,16]\n",
    "    x = model.encoder_embed(x)\n",
    "    #进入循环网络,得到记忆\n",
    "    #[b,8,16] -> [b,8,32],[1,b,32],[1,b,32]\n",
    "    out_x, (h, c) = model.encoder(x, None)\n",
    "\n",
    "    #[1,b,32],[1,b,32] -> [b,32],[b,32]\n",
    "    h = h.squeeze()\n",
    "    c = c.squeeze()\n",
    "\n",
    "    #初始化输入,每一个词的输入应该是上一个词的输出\n",
    "    #因为我们的y第一个词固定是<SOS>,所以直接以这个词开始\n",
    "    #[b]\n",
    "    out = torch.full((x.size(0), ), zidian['<SOS>'], dtype=torch.int64)\n",
    "    #[b] -> [b,16]\n",
    "    out = model.decoder_embed(out)\n",
    "\n",
    "    #循环生成9个词,收尾的两个标签没有预测的价值,直接忽略了\n",
    "    outs = []\n",
    "    for i in range(9):\n",
    "        \n",
    "        #[b,32] -> [b,1,32]\n",
    "        attn = model.get_attn(out_x,h)\n",
    "\n",
    "        #把每个词输入循环网络\n",
    "        #第一个词的记忆是x的最后一个词的记忆\n",
    "        #往后每个词的记忆是上一个词的记忆\n",
    "        #[b,16] -> [b,32],[b,32]\n",
    "        h, c = model.decoder_cell(out, (h, c))\n",
    "        \n",
    "        #[b,32],[b,32] -> [b,64]\n",
    "        attn = torch.cat([attn, h], dim=1)\n",
    "\n",
    "        #把每一步的记忆输出成词\n",
    "        #[b,64] -> [b,27]\n",
    "        out = model.out_fc(attn)\n",
    "\n",
    "        #把每一步的记忆输出成词\n",
    "        #[b,27] -> [b]\n",
    "        out = out.argmax(dim=1)\n",
    "        outs.append(out)\n",
    "\n",
    "        #把这一步的输出作为下一步的输入\n",
    "        #[b] -> [b,16]\n",
    "        out = model.decoder_embed(out)\n",
    "\n",
    "    #把所有的输出词组合成一句话\n",
    "    #[9,b]\n",
    "    outs = torch.stack(outs, dim=0)\n",
    "    #[9,b] -> [b,9]\n",
    "    outs = outs.permute(1, 0)\n",
    "\n",
    "    return outs\n",
    "\n",
    "\n",
    "#测试\n",
    "for i, data in enumerate(dataloader):\n",
    "    x, y = data\n",
    "    y_pred = predict(x)\n",
    "    for xi, yi, pi in zip(x, y, y_pred):\n",
    "        print(seq_to_str(xi), seq_to_str(yi), seq_to_str(pi))\n",
    "    break"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
